import torch
import yaml
import cv2
import argparse


import sys
sys.path.append("/home/daybeha/Documents/github/DeepLabV3_ws/src/superglue")
from Detectors import create_detector
from Matchers import create_matcher

from models.utils import AverageTimer
from utils.tools import *


class Matching(torch.nn.Module):
    """ Image Matching Frontend (SuperPoint + SuperGlue) """
    def __init__(self, config={}):
        super().__init__()
        # create detector
        self.detector = create_detector(config["detector"])
        # create matcher
        self.matcher = create_matcher(config["matcher"])


        # self.superpoint = SuperPoint(config.get('superpoint', {}))
        # self.superglue = SuperGlue(config.get('superglue', {}))

    def forward(self, data):
        """ Run SuperPoint (optionally) and SuperGlue
        SuperPoint is skipped if ['keypoints0', 'keypoints1'] exist in input
        Args:
          data: dictionary with minimal keys: ['image0', 'image1']
        """
        pred = {'ref': None, 'cur': None}

        # TODO 这块的显存存占用有待优化
        # Extract SuperPoint (keypoints, scores, descriptors) if not provided
        if 'keypoints0' not in data:
            # pred0 = self.superpoint({'image': data['image0']})
            # pred = {**pred, **{k+'0': v for k, v in pred0.items()}}
            pred['ref'] = self.detector(data['image0'])
            # pred0 = self.detector(data['image0'])
            # pred = {**pred, **{k+'0': v for k, v in pred0.items()}}
        if 'keypoints1' not in data:
            # pred1 = self.superpoint({'image': data['image1']})
            # pred = {**pred, **{k+'1': v for k, v in pred1.items()}}
            pred['cur'] = self.detector(data['image1'])
            # pred0 = self.detector(data['image1'])
            # pred = {**pred, **{k+'1': v for k, v in pred0.items()}}



        # Batch all features
        # We should either have i) one image per batch, or
        # ii) the same number of local features for all images in the batch.
        # data = {**data, **pred}
        #
        # for k in data:
        #     if isinstance(data[k], (list, tuple)):      # 将 list, tuple类转换为tensor  TODO 这部分或许可以优化
        #         data[k] = torch.stack(data[k])

        # Perform the matching
        # pred = {**pred, **self.superglue(data)}     # TODO pred这部分只用了matches0和matching_score0，matches1和matching_score1或可删去

        matches = self.matcher(pred)

        return matches

# 这一句至关重要！！！ 能节省至少一半显存！！！
torch.set_grad_enabled(False)

config = "/home/daybeha/Documents/github/DeepLabV3_ws/src/superglue/params/superpoint_supergluematch.yaml"
with open(config, 'r') as f:
    config = yaml.safe_load(f)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
matching = Matching(config).eval().to(device)

def compute_score(img0, img1):
    matches = matching({'image0': img0, 'image1': img1})
    score = np.mean(matches["match_score"].cpu().detach().numpy())
    return score


print(f"cuda avaliable: {torch.cuda.is_available()}")

def con_show(img0, img1):
    # 纵向连接 image = np.vstack((img0, img1))
    # 横向连接 image = np.concatenate([img0, img1], axis=1)
    image = np.concatenate((img0, img1))

    cv2.imshow("image in python", image)
    # cv2.imwrite("/home/tt/test/111.png", img)
    cv2.waitKey(0)


def show_img(img=None):
    if img is not None:
        cv2.imshow("image in python", img)
        # cv2.imwrite("/home/tt/test/111.png", img)
        cv2.waitKey(0)



if __name__ == "__main__":
    img0 = cv2.imread("/home/daybeha/Documents/Dataset/Kitti/sequences/00/image_0/000000.png")
    img1 = cv2.imread("/home/daybeha/Documents/Dataset/Kitti/sequences/00/image_0/000020.png")
    con_show(img0, img1)

    score = compute_score(img0, img1)
    print(f"score: {score}")
